import torch
import torchvision.models as models

def inference(model, data):
    with torch.no_grad():
        # warm up
        for _ in range(100):
            model(data)

        # measure
        import time
        measure_iter = 100
        start = time.time()
        for _ in range(measure_iter):
            output = model(data)
        end = time.time()

        duration = (end - start) * 1000
        latency = duration / measure_iter
        throughput = measure_iter / duration

        print("@hypertune {'name': 'latency (ms)'}")  # Add print statement of the form @hypertune {'name': str, 'higher_is_better': bool, 'target_val': int or float}`
        print(latency)  # Print the objective(s) you want to optimize. Make sure this is just an int or float to be minimzied or maximized.

def main(args):
    model = models.resnet50(pretrained=False)
    model.eval()

    data = torch.rand(1, 3, 224, 224)

    import intel_extension_for_pytorch as ipex

    if args.dtype == 'float32':
        model = ipex.optimize(model, dtype=torch.float32)
    elif args.dtype == 'bfloat16':
        model = ipex.optimize(model, dtype=torch.bfloat16)
    else:  # int8
        from intel_extension_for_pytorch.quantization import prepare, convert

        qconfig = ipex.quantization.default_static_qconfig
        model = prepare(model, qconfig, example_inputs=data, inplace=False)

        # calibration
        n_iter = 100
        for i in range(n_iter):
            model(data)

        model = convert(model)

    with torch.cpu.amp.autocast(enabled=args.dtype == 'bfloat16'):
        if args.torchscript:
            with torch.no_grad():
                model = torch.jit.trace(model, data)
                model = torch.jit.freeze(model)

        inference(model, data)

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--dtype', default='float32', choices=['float32', 'bfloat16', 'int8'])
    parser.add_argument("--torchscript", default=False, action="store_true")

    main(parser.parse_args())

print("Execution finished")
